package com.starhub.application.model.util;

import java.util.regex.Pattern;

/**
 * token计算工具类
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-07-15
 */
public class TokenUtil {
    
    /**
     * 中文字符正则表达式
     */
    private static final Pattern CHINESE_PATTERN = Pattern.compile("[\\u4e00-\\u9fa5]");
    
    /**
     * 英文单词正则表达式
     */
    private static final Pattern ENGLISH_WORD_PATTERN = Pattern.compile("[a-zA-Z]+");
    
    /**
     * 数字正则表达式
     */
    private static final Pattern NUMBER_PATTERN = Pattern.compile("\\d+");
    
    /**
     * 标点符号正则表达式
     */
    private static final Pattern PUNCTUATION_PATTERN = Pattern.compile("[\\p{P}\\p{Z}]");
    
    /**
     * 计算文本的token数量（基于字符数估算）
     * 适用于大多数大模型，如GPT系列
     * 
     * @param text 要计算的文本
     * @return token数量
     */
    public static int calculateTokensByChar(String text) {
        if (text == null || text.isEmpty()) {
            return 0;
        }
        
        // 中文字符通常1个字符对应1个token
        // 英文字符通常4个字符对应1个token
        int chineseCount = countChineseCharacters(text);
        int englishCount = text.length() - chineseCount;
        
        // 计算token数：中文字符数 + 英文字符数/4（向上取整）
        return chineseCount + (int) Math.ceil(englishCount / 4.0);
    }
    
    /**
     * 计算文本的token数量（基于词数估算）
     * 适用于需要更精确计算的场景
     * 
     * @param text 要计算的文本
     * @return token数量
     */
    public static int calculateTokensByWord(String text) {
        if (text == null || text.isEmpty()) {
            return 0;
        }
        
        int totalTokens = 0;
        
        // 计算中文字符
        int chineseCount = countChineseCharacters(text);
        totalTokens += chineseCount;
        
        // 计算英文单词
        int englishWordCount = countEnglishWords(text);
        totalTokens += englishWordCount;
        
        // 计算数字
        int numberCount = countNumbers(text);
        totalTokens += numberCount;
        
        // 计算标点符号和空格
        int punctuationCount = countPunctuation(text);
        totalTokens += punctuationCount;
        
        return totalTokens;
    }
    
    /**
     * 计算文本的token数量（GPT-3/4风格）
     * 基于GPT tokenizer的估算方法
     * 
     * @param text 要计算的文本
     * @return token数量
     */
    public static int calculateGPTTokens(String text) {
        if (text == null || text.isEmpty()) {
            return 0;
        }
        
        // GPT tokenizer的粗略估算：
        // 英文：1个token ≈ 0.75个单词
        // 中文：1个token ≈ 1个字符
        // 数字和标点：单独计算
        
        int chineseCount = countChineseCharacters(text);
        int englishWordCount = countEnglishWords(text);
        int numberCount = countNumbers(text);
        int punctuationCount = countPunctuation(text);
        
        // 英文单词转换为token（1个单词约等于1.33个token）
        int englishTokens = (int) Math.ceil(englishWordCount * 1.33);
        
        return chineseCount + englishTokens + numberCount + punctuationCount;
    }
    
    /**
     * 计算文本的token数量（Claude风格）
     * 基于Claude tokenizer的估算方法
     * 
     * @param text 要计算的文本
     * @return token数量
     */
    public static int calculateClaudeTokens(String text) {
        if (text == null || text.isEmpty()) {
            return 0;
        }
        
        // Claude tokenizer的粗略估算：
        // 英文：1个token ≈ 0.6个单词
        // 中文：1个token ≈ 1个字符
        // 数字和标点：单独计算
        
        int chineseCount = countChineseCharacters(text);
        int englishWordCount = countEnglishWords(text);
        int numberCount = countNumbers(text);
        int punctuationCount = countPunctuation(text);
        
        // 英文单词转换为token（1个单词约等于1.67个token）
        int englishTokens = (int) Math.ceil(englishWordCount * 1.67);
        
        return chineseCount + englishTokens + numberCount + punctuationCount;
    }
    
    /**
     * 统计中文字符数量
     * 
     * @param text 文本
     * @return 中文字符数量
     */
    private static int countChineseCharacters(String text) {
        return (int) CHINESE_PATTERN.matcher(text).results().count();
    }
    
    /**
     * 统计英文单词数量
     * 
     * @param text 文本
     * @return 英文单词数量
     */
    private static int countEnglishWords(String text) {
        return (int) ENGLISH_WORD_PATTERN.matcher(text).results().count();
    }
    
    /**
     * 统计数字数量
     * 
     * @param text 文本
     * @return 数字数量
     */
    private static int countNumbers(String text) {
        return (int) NUMBER_PATTERN.matcher(text).results().count();
    }
    
    /**
     * 统计标点符号和空格数量
     * 
     * @param text 文本
     * @return 标点符号和空格数量
     */
    private static int countPunctuation(String text) {
        return (int) PUNCTUATION_PATTERN.matcher(text).results().count();
    }
    
    /**
     * 获取文本的详细token分析
     * 
     * @param text 要分析的文本
     * @return 包含详细信息的TokenAnalysis对象
     */
    public static TokenAnalysis analyzeTokens(String text) {
        if (text == null || text.isEmpty()) {
            return new TokenAnalysis(0, 0, 0, 0, 0, 0);
        }
        
        int chineseCount = countChineseCharacters(text);
        int englishWordCount = countEnglishWords(text);
        int numberCount = countNumbers(text);
        int punctuationCount = countPunctuation(text);
        
        int charBasedTokens = calculateTokensByChar(text);
        int wordBasedTokens = calculateTokensByWord(text);
        
        return new TokenAnalysis(
            chineseCount,
            englishWordCount,
            numberCount,
            punctuationCount,
            charBasedTokens,
            wordBasedTokens
        );
    }
    
    /**
     * Token分析结果类
     */
    public static class TokenAnalysis {
        private final int chineseCharacters;
        private final int englishWords;
        private final int numbers;
        private final int punctuation;
        private final int charBasedTokens;
        private final int wordBasedTokens;
        
        public TokenAnalysis(int chineseCharacters, int englishWords, int numbers, 
                           int punctuation, int charBasedTokens, int wordBasedTokens) {
            this.chineseCharacters = chineseCharacters;
            this.englishWords = englishWords;
            this.numbers = numbers;
            this.punctuation = punctuation;
            this.charBasedTokens = charBasedTokens;
            this.wordBasedTokens = wordBasedTokens;
        }
        
        // Getters
        public int getChineseCharacters() { return chineseCharacters; }
        public int getEnglishWords() { return englishWords; }
        public int getNumbers() { return numbers; }
        public int getPunctuation() { return punctuation; }
        public int getCharBasedTokens() { return charBasedTokens; }
        public int getWordBasedTokens() { return wordBasedTokens; }
        
        @Override
        public String toString() {
            return String.format(
                "TokenAnalysis{chinese=%d, english=%d, numbers=%d, punctuation=%d, " +
                "charTokens=%d, wordTokens=%d}",
                chineseCharacters, englishWords, numbers, punctuation, 
                charBasedTokens, wordBasedTokens
            );
        }
    }
}
